import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.sequence import pad_sequences
import json

# 定义文件路径
CURRENT_DATA_PATH = 'current_user_data.json'

def load_new_user_data(file_path):
    """加载用户击键数据"""
    with open(file_path, 'r') as f:
        return json.load(f)

def predict_user(keystrokes, model, key_to_index, max_length=100):
    """预测用户身份"""
    sequence = []
    for i in range(len(keystrokes)):
        key_content = keystrokes[i]['key_content']
        press_time = keystrokes[i]['press_time']
        release_time = keystrokes[i]['release_time']
        duration = release_time - press_time
        
        if i > 0:
            interval = press_time - keystrokes[i-1]['press_time']
            flight_time = press_time - keystrokes[i-1]['release_time']
            prev_key_content = keystrokes[i-1]['key_content']
            acceleration = (duration - (keystrokes[i-1]['release_time'] - keystrokes[i-1]['press_time'])) / interval
        else:
            interval = flight_time = acceleration = 0.0
            prev_key_content = 'None'
        
        # 动态扩展key_to_index
        if key_content not in key_to_index:
            key_to_index[key_content] = len(key_to_index)
        if prev_key_content not in key_to_index:
            key_to_index[prev_key_content] = len(key_to_index)
        
        sequence.append([
            key_to_index[key_content],
            key_to_index[prev_key_content],
            press_time,
            release_time,
            duration,
            interval,
            flight_time,
            acceleration
        ])
    
    # 序列填充
    padded_sequence = pad_sequences([sequence], maxlen=max_length, padding='post', dtype='float32')
    predictions = model.predict(padded_sequence)
    return np.argmax(predictions), np.max(predictions)

def main():
    try:
        # 加载模型和配置文件
        model = load_model('model/lstm_model.h5')
        with open('model/lstm_key_to_index.json', 'r') as f:
            key_to_index = json.load(f)
        with open('model/lstm_label_to_index.json', 'r') as f:
            label_to_index = json.load(f)
        
        index_to_label = {v: k for k, v in label_to_index.items()}
        user_data = load_new_user_data(CURRENT_DATA_PATH)
        
        if not user_data:
            print("错误：没有有效的击键数据！")
            return
            
        for entry in user_data:
            subject = entry['subject']
            keystrokes = entry['keystrokes']
            
            if len(keystrokes) < 5:
                continue
                
            user_idx, confidence = predict_user(keystrokes, model, key_to_index)
            predicted_user = index_to_label.get(user_idx, "未知用户")
            
            print(f"预测用户: {predicted_user}, 置信度: {confidence:.2%}")
            
    except Exception as e:
        print(f"错误: {str(e)}")

if __name__ == "__main__":
    main()